import os
import re
import time
from typing import List

from sparkai.core.messages import ChatMessage

from iflytech_assistant.spark_client import spark, spark_streaming

from .role import Example, RoleMeta, SearchResult
from .role_play import generate as role_play_generate
from .avatar import avatar_generate


__all__ = ["generate_role_single_turn"]

TEMPLATE_FOLDER = os.path.join(os.path.dirname(__file__), "templates", "creator")

SEARCH_SYSTEM_PROMPT = """
你是一个人物信息检索器

人物角色的定义：影视娱乐明星、动漫、影视文学、戏剧、小说作品中的人物。

给定你一个输入内容，请你在网上搜索相关信息（主要参考百科数据）。

请先判断检索到的内容是否符合人物角色的定义，
- 如果符合人物角色的定义，请将该人物角色的信息进行总结并输出，总结可以简短一些，不超过50字。
- 如果不符合人物角色的定义，请输出“不符合人物角色的定义”。
- 如果没有搜到相关人物信息，请输出“未找到相关信息”。
"""

SEARCH_USER_PROMPT = """
输入内容：
```
{name}相关的人物背景信息
```
请在网上搜索相关信息，并判断总结。
"""


def web_search(name: str) -> SearchResult:
    search_prompt = SEARCH_USER_PROMPT.format(name=name)
    messages = [
        ChatMessage(role="system", content=SEARCH_SYSTEM_PROMPT),
        ChatMessage(
            role="user",
            content=search_prompt,
        ),
    ]
    try:
        response = spark.generate([messages], temperature=0.0)
        raw_text = response.generations[0][0].text
    except Exception:
        # Error Code: 10014, Error: 非常抱歉，根据相关法律法规，有关信息不予显示
        # 返回内容为空的情况
        return SearchResult(name=name)
    if "未找到" in raw_text or "不符合" in raw_text:
        return SearchResult(
            name=name,
            exists=False,
            description=raw_text,
            search_prompt=SEARCH_SYSTEM_PROMPT + "\n" + search_prompt,
        )
    else:
        return SearchResult(
            name=name,
            exists=True,
            description=raw_text,
            search_prompt=SEARCH_SYSTEM_PROMPT + "\n" + search_prompt,
        )


SYSTEM_PROMPT = open(
    os.path.join(TEMPLATE_FOLDER, "system_prompt.md"), encoding="utf-8"
).read()

SYSTEM_PROMPT_WITHOUT_EXAMPLES = open(
    os.path.join(TEMPLATE_FOLDER, "system_prompt_without_examples.md"), encoding="utf-8"
).read()


def build_system_prompt(generate_examples: bool = True) -> str:
    if generate_examples:
        return SYSTEM_PROMPT
    else:
        return SYSTEM_PROMPT_WITHOUT_EXAMPLES


USER_PROMPT = """
{background}
{lines}
{content}

请你按照指定json格式创建一个名为“{name}”的角色助手。

格式要求：
```json
{{
  "name": "角色名",
  "introduction": "角色介绍",
  "category": "所属分类专区，古风专区、聊天风格、人物角色、玩法这里面选择一个",
  "profile": "一句话描述"
}}
```
"""

BACKGROUNDS_PROMPT = """
背景信息: {background}
"""

LINES_PROMPT = """
角色曾经说过的话: {lines}
"""

CONTENT_PROMPT = """
content字段填入以下内容，:
{content}
rewrite字段填入根据角色设定改写润色的内容，不要回复，不要回复。
"""

APPAERANCE_PROMPT = """
角色名称：{name}
{introduction}

请根据角色名称，和用户提供的信息，生成一个人物外貌描述，不超过50字。
"""


def generate_appearance(name: str, introduction: str = "") -> str:
    prompt = APPAERANCE_PROMPT.format(name=name, introduction=introduction)

    messages = [
        ChatMessage(role="user", content=prompt),
    ]
    result = spark.generate([messages], temperature=0.3)
    raw_text = result.generations[0][0].text
    return raw_text


def build_user_prompt(
    name: str,
    background: str = "",
    lines: List[str] = [],
    content: List[str] = [],
) -> str:
    user_prompt = USER_PROMPT.format(
        name=name,
        background=(
            BACKGROUNDS_PROMPT.format(background=background) if background else ""
        ),
        lines=LINES_PROMPT.format(lines="".join(lines)) if lines else "",
        content=CONTENT_PROMPT.format(content="\n".join(content)) if content else "",
    )
    user_prompt = user_prompt.strip()
    user_prompt = re.sub(r"\n{3,}", "\n\n", user_prompt)
    return user_prompt


def generate_role_single_turn(
    name: str,
    lines: List[str] = [],
    examples: List[Example] = [],
    do_web_search: bool = False,
    generate_examples: bool = False,
    do_test: bool = False,
    test_user_input: List[str] = ["哈哈", "加油啊", "别生气了"],
    generate_avatar: bool = False,
) -> RoleMeta:
    start_time = time.time()
    if do_web_search:
        search_result: SearchResult = web_search(name)
    else:
        search_result: SearchResult = SearchResult(name=name, description="未调用搜索")
    search_time = time.time()
    system_prompt = build_system_prompt(generate_examples=generate_examples)
    if search_result.exists:
        user_prompt = build_user_prompt(
            name=name,
            background=search_result.description,
            lines=lines,
        )
    else:
        user_prompt = build_user_prompt(
            name=name,
            background="",
            lines=lines,
        )

    messages = [
        ChatMessage(role="system", content=system_prompt.strip()),
        ChatMessage(
            role="user",
            content=user_prompt.strip(),
        ),
    ]
    response = spark_streaming.stream(messages, temperature=0.0)
    raw_text = ""
    prefilling_time = time.time()
    for i, message in enumerate(response):
        if i == 1:
            prefilling_time = time.time()
        raw_text += message.content

    end_time = time.time()
    role = RoleMeta.from_str(raw_text)
    role.search_result = search_result
    role.lines = lines
    # 显式传递examples，则覆盖自动生成的
    if examples:
        for example in examples:
            role.content.append(example.user_input)
            role.rewrite.append(example.role_output)
    role.search_time = search_time - start_time
    role.prefilling_time = prefilling_time - search_time
    role.generation_time = end_time - prefilling_time

    if do_test:
        for user_input in test_user_input:
            example = role_play_generate(role, user_input)
            role.examples.append(example)

    # generate appearance
    role.appearance = generate_appearance(name, role.introduction)
    if generate_avatar:
        role.avatar = avatar_generate(f"{role.name}: {role.introduction}")
    return role
